/*
  PubSubClient.cpp - A simple client for MQTT.
  Nick O'Leary
  http://knolleary.net
*/

#include "PubSubClient.h"

PubSubClient::PubSubClient()
{
    this->_state = MQTT_DISCONNECTED;
    this->_client = NULL;
    this->stream = NULL;
    setCallback(NULL);
}

PubSubClient::PubSubClient(Client &client)
{
    this->_state = MQTT_DISCONNECTED;
    setClient(client);
    this->stream = NULL;
}

PubSubClient::PubSubClient(IPAddress addr, uint16_t port, Client &client)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(addr, port);
    setClient(client);
    this->stream = NULL;
}
PubSubClient::PubSubClient(IPAddress addr, uint16_t port, Client &client, Stream &stream)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(addr, port);
    setClient(client);
    setStream(stream);
}
PubSubClient::PubSubClient(IPAddress addr, uint16_t port, MQTT_CALLBACK_SIGNATURE, Client &client)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(addr, port);
    setCallback(callback);
    setClient(client);
    this->stream = NULL;
}
PubSubClient::PubSubClient(IPAddress addr, uint16_t port, MQTT_CALLBACK_SIGNATURE, Client &client, Stream &stream)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(addr, port);
    setCallback(callback);
    setClient(client);
    setStream(stream);
}

PubSubClient::PubSubClient(uint8_t *ip, uint16_t port, Client &client)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(ip, port);
    setClient(client);
    this->stream = NULL;
}
PubSubClient::PubSubClient(uint8_t *ip, uint16_t port, Client &client, Stream &stream)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(ip, port);
    setClient(client);
    setStream(stream);
}
PubSubClient::PubSubClient(uint8_t *ip, uint16_t port, MQTT_CALLBACK_SIGNATURE, Client &client)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(ip, port);
    setCallback(callback);
    setClient(client);
    this->stream = NULL;
}
PubSubClient::PubSubClient(uint8_t *ip, uint16_t port, MQTT_CALLBACK_SIGNATURE, Client &client, Stream &stream)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(ip, port);
    setCallback(callback);
    setClient(client);
    setStream(stream);
}

PubSubClient::PubSubClient(const char *domain, uint16_t port, Client &client)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(domain, port);
    setClient(client);
    this->stream = NULL;
}
PubSubClient::PubSubClient(const char *domain, uint16_t port, Client &client, Stream &stream)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(domain, port);
    setClient(client);
    setStream(stream);
}
PubSubClient::PubSubClient(const char *domain, uint16_t port, MQTT_CALLBACK_SIGNATURE, Client &client)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(domain, port);
    setCallback(callback);
    setClient(client);
    this->stream = NULL;
}
PubSubClient::PubSubClient(const char *domain, uint16_t port, MQTT_CALLBACK_SIGNATURE, Client &client, Stream &stream)
{
    this->_state = MQTT_DISCONNECTED;
    setServer(domain, port);
    setCallback(callback);
    setClient(client);
    setStream(stream);
}

boolean PubSubClient::connect(const char *id)
{
    return connect(id, NULL, NULL, 0, 0, 0, 0, 1);
}

boolean PubSubClient::connect(const char *id, const char *user, const char *pass)
{
    return connect(id, user, pass, 0, 0, 0, 0, 1);
}

boolean PubSubClient::connect(const char *id, const char *willTopic, uint8_t willQos, boolean willRetain, const char *willMessage)
{
    return connect(id, NULL, NULL, willTopic, willQos, willRetain, willMessage, 1);
}

boolean PubSubClient::connect(const char *id, const char *user, const char *pass, const char *willTopic, uint8_t willQos, boolean willRetain, const char *willMessage)
{
    return connect(id, user, pass, willTopic, willQos, willRetain, willMessage, 1);
}

boolean PubSubClient::connect(const char *id, const char *user, const char *pass, const char *willTopic, uint8_t willQos, boolean willRetain, const char *willMessage, boolean cleanSession)
{
    if (!connected())
    {
        int result = 0;

        if (domain != NULL)
        {
            result = _client->connect(this->domain, this->port);
        }
        else
        {
            result = _client->connect(this->ip, this->port);
        }
        if (result == 1)
        {
            nextMsgId = 1;
            // Leave room in the buffer for header and variable length field
            uint16_t length = MQTT_MAX_HEADER_SIZE;
            unsigned int j;

#if MQTT_VERSION == MQTT_VERSION_3_1
            uint8_t d[9] = {0x00, 0x06, 'M', 'Q', 'I', 's', 'd', 'p', MQTT_VERSION};
#define MQTT_HEADER_VERSION_LENGTH 9
#elif MQTT_VERSION == MQTT_VERSION_3_1_1
            uint8_t d[7] = {0x00, 0x04, 'M', 'Q', 'T', 'T', MQTT_VERSION};
#define MQTT_HEADER_VERSION_LENGTH 7
#endif
            for (j = 0; j < MQTT_HEADER_VERSION_LENGTH; j++)
            {
                buffer[length++] = d[j];
            }

            uint8_t v;
            if (willTopic)
            {
                v = 0x04 | (willQos << 3) | (willRetain << 5);
            }
            else
            {
                v = 0x00;
            }
            if (cleanSession)
            {
                v = v | 0x02;
            }

            if(user != NULL)
            {
                v = v | 0x80;

                if(pass != NULL)
                {
                    v = v | (0x80 >> 1);
                }
            }

            buffer[length++] = v;

            buffer[length++] = ((MQTT_KEEPALIVE) >> 8);
            buffer[length++] = ((MQTT_KEEPALIVE) & 0xFF);

            CHECK_STRING_LENGTH(length, id)
            length = writeString(id, buffer, length);
            if (willTopic)
            {
                CHECK_STRING_LENGTH(length, willTopic)
                length = writeString(willTopic, buffer, length);
                CHECK_STRING_LENGTH(length, willMessage)
                length = writeString(willMessage, buffer, length);
            }

            if(user != NULL)
            {
                CHECK_STRING_LENGTH(length, user)
                length = writeString(user, buffer, length);
                if(pass != NULL)
                {
                    CHECK_STRING_LENGTH(length, pass)
                    length = writeString(pass, buffer, length);
                }
            }

            write(MQTTCONNECT, buffer, length - MQTT_MAX_HEADER_SIZE);

            lastInActivity = lastOutActivity = millis();

            while (!_client->available())
            {
                unsigned long t = millis();
                if (t - lastInActivity >= ((int32_t) MQTT_SOCKET_TIMEOUT * 1000UL))
                {
                    _state = MQTT_CONNECTION_TIMEOUT;
                    _client->stop();
                    return false;
                }
            }
            uint8_t llen;
            uint16_t len = readPacket(&llen);

            if (len == 4)
            {
                if (buffer[3] == 0)
                {
                    lastInActivity = millis();
                    pingOutstanding = false;
                    _state = MQTT_CONNECTED;
                    return true;
                }
                else
                {
                    _state = buffer[3];
                }
            }
            _client->stop();
        }
        else
        {
            _state = MQTT_CONNECT_FAILED;
        }
        return false;
    }
    return true;
}

// reads a byte into result
boolean PubSubClient::readByte(uint8_t *result)
{
    uint32_t previousMillis = millis();
    while(!_client->available())
    {
        //     yield();
        uint32_t currentMillis = millis();
        if(currentMillis - previousMillis >= ((int32_t) MQTT_SOCKET_TIMEOUT * 1000))
        {
            return false;
        }
    }
    *result = _client->read();
    return true;
}

// reads a byte into result[*index] and increments index
boolean PubSubClient::readByte(uint8_t *result, uint16_t *index)
{
    uint16_t current_index = *index;
    uint8_t *write_address = &(result[current_index]);
    if(readByte(write_address))
    {
        *index = current_index + 1;
        return true;
    }
    return false;
}

uint16_t PubSubClient::readPacket(uint8_t *lengthLength)
{
    uint16_t len = 0;
    if(!readByte(buffer, &len)) return 0;
    bool isPublish = (buffer[0] & 0xF0) == MQTTPUBLISH;
    uint32_t multiplier = 1;
    uint16_t length = 0;
    uint8_t digit = 0;
    uint16_t skip = 0;
    uint8_t start = 0;

    do
    {
        if (len == 5)
        {
            // Invalid remaining length encoding - kill the connection
            _state = MQTT_DISCONNECTED;
            _client->stop();
            return 0;
        }
        if(!readByte(&digit)) return 0;
        buffer[len++] = digit;
        length += (digit & 127) * multiplier;
        multiplier *= 128;
    }
    while ((digit & 128) != 0);
    *lengthLength = len - 1;

    if (isPublish)
    {
        // Read in topic length to calculate bytes to skip over for Stream writing
        if(!readByte(buffer, &len)) return 0;
        if(!readByte(buffer, &len)) return 0;
        skip = (buffer[*lengthLength + 1] << 8) + buffer[*lengthLength + 2];
        start = 2;
        if (buffer[0]&MQTTQOS1)
        {
            // skip message id
            skip += 2;
        }
    }

    for (uint16_t i = start; i < length; i++)
    {
        if(!readByte(&digit)) return 0;
        if (this->stream)
        {
            if (isPublish && len - *lengthLength - 2 > skip)
            {
                this->stream->write(digit);
            }
        }
        if (len < MQTT_MAX_PACKET_SIZE)
        {
            buffer[len] = digit;
        }
        len++;
    }

    if (!this->stream && len > MQTT_MAX_PACKET_SIZE)
    {
        len = 0; // This will cause the packet to be ignored.
    }

    return len;
}

boolean PubSubClient::loop()
{
    if (connected())
    {
        unsigned long t = millis();
        if ((t - lastInActivity > MQTT_KEEPALIVE * 1000UL) || (t - lastOutActivity > MQTT_KEEPALIVE * 1000UL))
        {
            if (pingOutstanding)
            {
                this->_state = MQTT_CONNECTION_TIMEOUT;
                _client->stop();
                return false;
            }
            else
            {
                buffer[0] = MQTTPINGREQ;
                buffer[1] = 0;
                _client->write(buffer, 2);
                lastOutActivity = t;
                lastInActivity = t;
                pingOutstanding = true;
            }
        }
        if (_client->available())
        {
            uint8_t llen;
            uint16_t len = readPacket(&llen);
            uint16_t msgId = 0;
            uint8_t *payload;
            if (len > 0)
            {
                lastInActivity = t;
                uint8_t type = buffer[0] & 0xF0;
                if (type == MQTTPUBLISH)
                {
                    if (callback)
                    {
                        uint16_t tl = (buffer[llen + 1] << 8) + buffer[llen + 2]; /* topic length in bytes */
                        memmove(buffer + llen + 2, buffer + llen + 3, tl); /* move topic inside buffer 1 byte to front */
                        buffer[llen + 2 + tl] = 0; /* end the topic as a 'C' string with \x00 */
                        char *topic = (char *) buffer + llen + 2;
                        // msgId only present for QOS>0
                        if ((buffer[0] & 0x06) == MQTTQOS1)
                        {
                            msgId = (buffer[llen + 3 + tl] << 8) + buffer[llen + 3 + tl + 1];
                            payload = buffer + llen + 3 + tl + 2;
                            callback(topic, payload, len - llen - 3 - tl - 2);

                            buffer[0] = MQTTPUBACK;
                            buffer[1] = 2;
                            buffer[2] = (msgId >> 8);
                            buffer[3] = (msgId & 0xFF);
                            _client->write(buffer, 4);
                            lastOutActivity = t;

                        }
                        else
                        {
                            payload = buffer + llen + 3 + tl;
                            callback(topic, payload, len - llen - 3 - tl);
                        }
                    }
                }
                else if (type == MQTTPINGREQ)
                {
                    buffer[0] = MQTTPINGRESP;
                    buffer[1] = 0;
                    _client->write(buffer, 2);
                }
                else if (type == MQTTPINGRESP)
                {
                    pingOutstanding = false;
                }
            }
            else if (!connected())
            {
                // readPacket has closed the connection
                return false;
            }
        }
        return true;
    }
    return false;
}

boolean PubSubClient::publish(const char *topic, const char *payload)
{
    return publish(topic, (const uint8_t *)payload, strlen(payload), false);
}

boolean PubSubClient::publish(const char *topic, const char *payload, boolean retained)
{
    return publish(topic, (const uint8_t *)payload, strlen(payload), retained);
}

boolean PubSubClient::publish(const char *topic, const uint8_t *payload, unsigned int plength)
{
    return publish(topic, payload, plength, false);
}

boolean PubSubClient::publish(const char *topic, const uint8_t *payload, unsigned int plength, boolean retained)
{
    if (connected())
    {
        if (MQTT_MAX_PACKET_SIZE < MQTT_MAX_HEADER_SIZE + 2 + strlen(topic) + plength)
        {
            // Too long
            return false;
        }
        // Leave room in the buffer for header and variable length field
        uint16_t length = MQTT_MAX_HEADER_SIZE;
        length = writeString(topic, buffer, length);
        uint16_t i;
        for (i = 0; i < plength; i++)
        {
            buffer[length++] = payload[i];
        }
        uint8_t header = MQTTPUBLISH;
        if (retained)
        {
            header |= 1;
        }
        return write(header, buffer, length - MQTT_MAX_HEADER_SIZE);
    }
    return false;
}

boolean PubSubClient::publish_P(const char *topic, const char *payload, boolean retained)
{
    return publish_P(topic, (const uint8_t *)payload, strlen(payload), retained);
}

boolean PubSubClient::publish_P(const char *topic, const uint8_t *payload, unsigned int plength, boolean retained)
{
    uint8_t llen = 0;
    uint8_t digit;
    unsigned int rc = 0;
    uint16_t tlen;
    unsigned int pos = 0;
    unsigned int i;
    uint8_t header;
    unsigned int len;

    if (!connected())
    {
        return false;
    }

    tlen = strlen(topic);

    header = MQTTPUBLISH;
    if (retained)
    {
        header |= 1;
    }
    buffer[pos++] = header;
    len = plength + 2 + tlen;
    do
    {
        digit = len % 128;
        len = len / 128;
        if (len > 0)
        {
            digit |= 0x80;
        }
        buffer[pos++] = digit;
        llen++;
    }
    while(len > 0);

    pos = writeString(topic, buffer, pos);

    rc += _client->write(buffer, pos);

    for (i = 0; i < plength; i++)
    {
        rc += _client->write(*(payload + i));
    }

    lastOutActivity = millis();

    return rc == tlen + 4 + plength;
}

boolean PubSubClient::beginPublish(const char *topic, unsigned int plength, boolean retained)
{
    if (connected())
    {
        // Send the header and variable length field
        uint16_t length = MQTT_MAX_HEADER_SIZE;
        length = writeString(topic, buffer, length);
        uint16_t i;
        uint8_t header = MQTTPUBLISH;
        if (retained)
        {
            header |= 1;
        }
        size_t hlen = buildHeader(header, buffer, plength + length - MQTT_MAX_HEADER_SIZE);
        uint16_t rc = _client->write(buffer + (MQTT_MAX_HEADER_SIZE - hlen), length - (MQTT_MAX_HEADER_SIZE - hlen));
        lastOutActivity = millis();
        return (rc == (length - (MQTT_MAX_HEADER_SIZE - hlen)));
    }
    return false;
}

int PubSubClient::endPublish()
{
    return 1;
}

size_t PubSubClient::write(uint8_t data)
{
    lastOutActivity = millis();
    return _client->write(data);
}

size_t PubSubClient::write(const uint8_t *buffer, size_t size)
{
    lastOutActivity = millis();
    return _client->write(buffer, size);
}

size_t PubSubClient::buildHeader(uint8_t header, uint8_t *buf, uint16_t length)
{
    uint8_t lenBuf[4];
    uint8_t llen = 0;
    uint8_t digit;
    uint8_t pos = 0;
    uint16_t len = length;
    do
    {
        digit = len % 128;
        len = len / 128;
        if (len > 0)
        {
            digit |= 0x80;
        }
        lenBuf[pos++] = digit;
        llen++;
    }
    while(len > 0);

    buf[4 - llen] = header;
    for (int i = 0; i < llen; i++)
    {
        buf[MQTT_MAX_HEADER_SIZE - llen + i] = lenBuf[i];
    }
    return llen + 1; // Full header size is variable length bit plus the 1-byte fixed header
}

boolean PubSubClient::write(uint8_t header, uint8_t *buf, uint16_t length)
{
    uint16_t rc;
    uint8_t hlen = buildHeader(header, buf, length);

#ifdef MQTT_MAX_TRANSFER_SIZE
    uint8_t *writeBuf = buf + (MQTT_MAX_HEADER_SIZE - hlen);
    uint16_t bytesRemaining = length + hlen; //Match the length type
    uint8_t bytesToWrite;
    boolean result = true;
    while((bytesRemaining > 0) && result)
    {
        bytesToWrite = (bytesRemaining > MQTT_MAX_TRANSFER_SIZE) ? MQTT_MAX_TRANSFER_SIZE : bytesRemaining;
        rc = _client->write(writeBuf, bytesToWrite);
        result = (rc == bytesToWrite);
        bytesRemaining -= rc;
        writeBuf += rc;
    }
    return result;
#else
    rc = _client->write(buf + (MQTT_MAX_HEADER_SIZE - hlen), length + hlen);
    lastOutActivity = millis();
    return (rc == hlen + length);
#endif
}

boolean PubSubClient::subscribe(const char *topic)
{
    return subscribe(topic, 0);
}

boolean PubSubClient::subscribe(const char *topic, uint8_t qos)
{
    if (qos > 1)
    {
        return false;
    }
    if (MQTT_MAX_PACKET_SIZE < 9 + strlen(topic))
    {
        // Too long
        return false;
    }
    if (connected())
    {
        // Leave room in the buffer for header and variable length field
        uint16_t length = MQTT_MAX_HEADER_SIZE;
        nextMsgId++;
        if (nextMsgId == 0)
        {
            nextMsgId = 1;
        }
        buffer[length++] = (nextMsgId >> 8);
        buffer[length++] = (nextMsgId & 0xFF);
        length = writeString((char *)topic, buffer, length);
        buffer[length++] = qos;
        return write(MQTTSUBSCRIBE | MQTTQOS1, buffer, length - MQTT_MAX_HEADER_SIZE);
    }
    return false;
}

boolean PubSubClient::unsubscribe(const char *topic)
{
    if (MQTT_MAX_PACKET_SIZE < 9 + strlen(topic))
    {
        // Too long
        return false;
    }
    if (connected())
    {
        uint16_t length = MQTT_MAX_HEADER_SIZE;
        nextMsgId++;
        if (nextMsgId == 0)
        {
            nextMsgId = 1;
        }
        buffer[length++] = (nextMsgId >> 8);
        buffer[length++] = (nextMsgId & 0xFF);
        length = writeString(topic, buffer, length);
        return write(MQTTUNSUBSCRIBE | MQTTQOS1, buffer, length - MQTT_MAX_HEADER_SIZE);
    }
    return false;
}

void PubSubClient::disconnect()
{
    buffer[0] = MQTTDISCONNECT;
    buffer[1] = 0;
    _client->write(buffer, 2);
    _state = MQTT_DISCONNECTED;
    _client->flush();
    _client->stop();
    lastInActivity = lastOutActivity = millis();
}

uint16_t PubSubClient::writeString(const char *string, uint8_t *buf, uint16_t pos)
{
    const char *idp = string;
    uint16_t i = 0;
    pos += 2;
    while (*idp)
    {
        buf[pos++] = *idp++;
        i++;
    }
    buf[pos - i - 2] = (i >> 8);
    buf[pos - i - 1] = (i & 0xFF);
    return pos;
}


boolean PubSubClient::connected()
{
    boolean rc;
    if (_client == NULL )
    {
        rc = false;
    }
    else
    {
        rc = (int)_client->connected();
        if (!rc)
        {
            if (this->_state == MQTT_CONNECTED)
            {
                this->_state = MQTT_CONNECTION_LOST;
                _client->flush();
                _client->stop();
            }
        }
    }
    return rc;
}

PubSubClient &PubSubClient::setServer(uint8_t *ip, uint16_t port)
{
    IPAddress addr(ip[0], ip[1], ip[2], ip[3]);
    return setServer(addr, port);
}

PubSubClient &PubSubClient::setServer(IPAddress ip, uint16_t port)
{
    this->ip = ip;
    this->port = port;
    this->domain = NULL;
    return *this;
}

PubSubClient &PubSubClient::setServer(const char *domain, uint16_t port)
{
    this->domain = domain;
    this->port = port;
    return *this;
}

PubSubClient &PubSubClient::setCallback(MQTT_CALLBACK_SIGNATURE)
{
    this->callback = callback;
    return *this;
}

PubSubClient &PubSubClient::setClient(Client &client)
{
    this->_client = &client;
    return *this;
}

PubSubClient &PubSubClient::setStream(Stream &stream)
{
    this->stream = &stream;
    return *this;
}

int PubSubClient::state()
{
    return this->_state;
}
